{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  HMinCaml.Closure
-- Copyright   :  (c) Stephen Tetley 2012
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Closure conversion.
--
--------------------------------------------------------------------------------

module HMinCaml.Closure
  (

    Closure(..)
  , Expr(..)
  , Fundef(..)
  , Prog

  , ppExpr

  , freevars

  , closure

  ) where

import qualified HMinCaml.HMonad as H
import HMinCaml.Id
import qualified HMinCaml.KNormal as K
import qualified HMinCaml.Type as T
import HMinCaml.Typing ( ExtEnv, extEnv, lookupExtEnv )

import Text.PrettyPrint.HughesPJ.PrettyExpr.Ocaml   -- package: hpj-pretty-expr


import Text.PrettyPrint.HughesPJ

import Control.Applicative
import Control.Monad
-- import Data.List ( intersperse )
import qualified Data.Map as M
import qualified Data.Set as S


data Closure = Closure { entry :: IdL, actual_freevars :: [IdS] }
  deriving (Eq,Show)

data Expr = Unit 
          | Int Int
          | Float Double
          | Neg IdS
          | Add IdS IdS
          | Sub IdS IdS
          | FNeg IdS
          | FAdd IdS IdS
          | FSub IdS IdS
          | FMul IdS IdS
          | FDiv IdS IdS
          | IfEq IdS IdS Expr Expr
          | IfLE IdS IdS Expr Expr
          | Let (IdS,T.Type) Expr Expr
          | Var IdS
          | MakeCls (IdS,T.Type) Closure Expr
          | AppCls IdS [IdS]
          | AppDir IdL [IdS]
          | Tuple [IdS]
          | LetTuple [(IdS, T.Type)] IdS Expr
          | Get IdS IdS
          | Put IdS IdS IdS
          | ExtArray IdL
  deriving (Eq,Show)

-- Fundef is the same as Syntax, except for Expr...

data Fundef = Fundef 
    { fun_head       :: (IdL, T.Type)
    , fun_args       :: [(IdS, T.Type)]
    , fun_formal_fvs :: [(IdS, T.Type)]
    , fun_body       :: Expr
    } 
  deriving (Eq, Show)


data Prog = Prog [Fundef] Expr
  deriving (Eq,Show)



ppExpr :: Expr -> Doc
ppExpr = unparse . exprDoc



exprDoc :: Expr -> DocE
exprDoc Unit                 = literal "()"
exprDoc (Int i)              = Atom $ int i 
exprDoc (Float d)            = Atom $ double $ realToFrac d 
exprDoc (Neg x)              = inegateU $ literal x 
exprDoc (Add x y)            = iaddB (literal x) (literal y)
exprDoc (Sub x y)            = isubtractB (literal x) (literal y)
exprDoc (FNeg x)             = fnegateU $ literal x 
exprDoc (FAdd x y)           = faddB (literal x) (literal y)
exprDoc (FSub x y)           = fsubtractB (literal x) (literal y)
exprDoc (FMul x y)           = fmultiplyB (literal x) (literal y)
exprDoc (FDiv x y)           = fdivideB (literal x) (literal y)
exprDoc (IfEq x y e1 e2)     = 
    ppIf (structuralEqB (literal x) (literal y)) e1 e2

exprDoc (IfLE x y e1 e2)     = 
    ppIf (lessThanEqB (literal x) (literal y)) e1 e2


exprDoc (Let (v,t) e b)      = ppLet v t e b

exprDoc (Var s)              = literal s
{-
exprDoc (LetRec fd e)        = 
    Atom $ ppFundef fd <+> text "in" <+> ppExpr e

exprDoc (App x xs)           = funAppB (literal x) (map literal xs)

exprDoc (Tuple xs)           = 
    Atom $ parens $ hcat $ intersperse comma $ map text xs

exprDoc (LetTuple xts x e)   = 
    Atom $ text "let" <+> ppTupleVars xts <+> char '=' <+> text x
                      <+> text "in" <+> ppExpr e
-}
exprDoc (Get x y)          = 
    Atom $ text x <> char '.' <> parens (text y)

exprDoc (Put x y z)        = 
    Atom $ text x <> char '.' <> parens (text y) 
                  <+> text "<-" <+> text z

exprDoc (ExtArray x)       = 
    Atom $ text "ext_array" <+> text (show x)

ppLet :: String -> T.Type -> Expr -> Expr -> DocE
ppLet v ty e b = 
    Atom $ text "let" <+> text v <+> text "::" <+> T.ppType ty
                      <+> char '=' <+> ppExpr e <+> text "in" 
                      $+$ nest 2 (ppExpr b)
{-
ppFundef :: Fundef -> Doc
ppFundef (Fundef hd args e) = 
    text "let rec" <+> text (fst hd) <+> hsep (map (text .fst) args) <+> char '=' 
      $+$ nest 2 (ppExpr e)

ppTupleVars :: [(IdS, T.Type)] -> Doc
ppTupleVars = parens . hcat . intersperse comma . map (text . fst)
-}

ppIf :: DocE -> Expr -> Expr -> DocE
ppIf clause e1 e2 = 
    Atom $ text "if" <+> unparse clause 
                     $+$ nest 2 (text "then" <+> ppExpr e1)
                     $+$ nest 2 (text "else" <+> ppExpr e2)


--------------------------------------------------------------------------------
-- Free variables

freevars :: Expr -> S.Set IdS
freevars Unit                 = S.empty
freevars (Int {})             = S.empty
freevars (Float {})           = S.empty
freevars (Neg x)              = S.singleton x
freevars (Add x y)            = S.fromList [x,y]
freevars (Sub x y)            = S.fromList [x,y]
freevars (FNeg x)             = S.singleton x
freevars (FAdd x y)           = S.fromList [x,y]
freevars (FSub x y)           = S.fromList [x,y]
freevars (FMul x y)           = S.fromList [x,y]
freevars (FDiv x y)           = S.fromList [x,y]
freevars (IfEq x y e1 e2)     = 
    S.insert x $ S.insert y $ S.union (freevars e1) (freevars e2)

freevars (IfLE x y e1 e2)     = 
    S.insert x $ S.insert y $ S.union (freevars e1) (freevars e2)

freevars (Let (x,_) e1 e2)    =
    S.union (freevars e1) (S.delete x $ freevars e2)

freevars (Var x)              = S.singleton x

freevars (MakeCls (x,t) (Closure { entry = l, actual_freevars = ys}) e) = 
    S.delete x $ S.union (S.fromList ys) (freevars e)

freevars (AppCls x xs)        = S.fromList (x:xs)
freevars (AppDir _ xs)        = S.fromList xs
freevars (Tuple xs)           = S.fromList xs
freevars (LetTuple xts y e)   = 
    S.insert y (freevars e S.\\ S.fromList (map fst xts))

freevars (Get x y)            = S.fromList [x,y]
freevars (Put x y z)          = S.fromList [x,y,z]
freevars (ExtArray _)         = S.empty




--------------------------------------------------------------------------------
-- k normalization

-- Need to also account for the typing env...

-- Monad needs state or writer...
type Mon a = H.HMonad a

type Env = M.Map IdS T.Type  

type Known = S.Set IdS


closure :: K.Expr -> H.HMonad Expr
closure expr = cc M.empty S.empty expr

cc :: Env -> Known -> K.Expr -> Mon Expr
cc _   _   (K.Unit)                 = return Unit
cc _   _   (K.Int i)                = return $ Int i
cc _   _   (K.Float d)              = return $ Float d
cc _   _   (K.Neg x)                = return $ Neg x
cc _   _   (K.Add x y)              = return $ Add x y
cc _   _   (K.Sub x y)              = return $ Sub x y
cc _   _   (K.FNeg x)               = return $ FNeg x
cc _   _   (K.FAdd x y)             = return $ FAdd x y
cc _   _   (K.FSub x y)             = return $ FSub x y
cc _   _   (K.FMul x y)             = return $ FMul x y
cc _   _   (K.FDiv x y)             = return $ FDiv x y
cc env kwn (K.IfEq x y e1 e2)       = 
    IfEq x y <$> cc env kwn e1 <*> cc env kwn e2

cc env kwn (K.IfLE x y e1 e2)       = 
    IfLE x y <$> cc env kwn e1 <*> cc env kwn e2



